import json
from django.core.management import BaseCommand
from django.conf import settings
from kafka import KafkaProducer, KafkaConsumer

KAFKA_PROD_ADDRESS = ["172.20.3.233:9092", "172.20.2.146:9092", "172.20.1.215:9092"]
KAFKA_TEST_ADDRESS = ['3.1.242.202:9092']
DEFAULT_SYNC_RATIO = 100
SYNC_RATIO_CONFIG = {
    'agent': 1,
    'mf_pay': 1
}


class Command(BaseCommand):
    def add_arguments(self, parser):
        parser.add_argument('--tenant_name', dest='tenant_name', help='loki|witch|zs')

    @staticmethod
    def get_compatible_event_data(message):
        compatible_target = ['aid', 'ip', 'cvc', 'chn', 'user_id', 's2s_appid', 'device_type']
        compatible_event_data = {}
        for key, value in message.items():
            if key in compatible_target:
                compatible_event_data['_' + key] = value
            else:
                compatible_event_data[key] = value
        message = compatible_event_data
        return message

    def handle(self, **kwargs):
        tenant_name = kwargs['tenant_name']
        assert tenant_name in settings.CP_INDEX + settings.DWC_INDEX + settings.AGENT_INDEX + settings.MF_PAY_INDEX
        consumer = KafkaConsumer(tenant_name, group_id='kafka_prod_to_test_' + tenant_name,
                                 bootstrap_servers=KAFKA_PROD_ADDRESS,
                                 value_deserializer=lambda m: json.loads(m.decode('utf-8')))
        producer = KafkaProducer(bootstrap_servers=KAFKA_TEST_ADDRESS,
                                 value_serializer=lambda x: json.dumps(x).encode('utf-8'))
        while True:
            batch_size = 0
            for each in consumer:
                try:
                    message = json.loads(each.value['message']) if not isinstance(
                        each.value['message'], dict) else each.value['message']
                except:
                    continue
                message = self.get_compatible_event_data(message)
                batch_size += 1
                sync_ratio = SYNC_RATIO_CONFIG.get(tenant_name, DEFAULT_SYNC_RATIO)
                if batch_size % sync_ratio == 0 or message['_event_id'] == 'register':
                    producer.send(tenant_name, each.value)
                    if message['_event_id'] != 'register':
                        batch_size = 0

